from imap_engine import EngineIMAP
from optimizer import Optimizer
import numpy as np
import sys
import os
import time
import matplotlib.pyplot as plt	

class BOiLS(object):
    def __init__(self, input_file, output_file=None) -> None:
        self.actions = ['blance','rewrite','rewrite -z','refactor','refactor -z']
        self.alphabet = len(self.actions)
        self.input_file = input_file
        
        if output_file == None:
            output_file = input_file+'seq'
        self.engine = EngineIMAP(input_file, output_file)
        
        #所有观测点
        self.Sample_X = []
        self.fX_area = []
        self.fX_level = []
        self.fX_obj = [] # area + level or QoR
        
    def run(self,n_tol_eval = 50, n_init = 20, seq_length = 7, gp_train_steps=300):
        self.seq_length = seq_length
        self.gp_train_steps = gp_train_steps
        #用于循环搜索的优化器,optim只负责给出序列,取完init的点后才有gp,此时已经有y了,optim不需要imap或yosys
        optim = Optimizer(seq_length=self.seq_length, alphabet=self.alphabet, n_init=n_init, gp_train_steps=self.gp_train_steps)
        
        for i in range(n_tol_eval):
            x_next = optim.suggest()
            print("{:2d}/{} suggest{}".format(i,n_tol_eval,x_next))
            y_next = self.evaluate(x_next)
            optim.observe(x_next,y_next)
            
        self.get_best()
        
    def evaluate(self,x):
        area, level, qor = self.engine.QoR(x, detail=True)
        self.Sample_X.append(x)
        self.fX_area.append(area)
        self.fX_level.append(level)
        self.fX_obj.append(qor)
        return qor    
    
    def get_best(self):
        index = np.argmin(self.fX_obj)
        print("best_seq:",self.Sample_X[index], " , index:", index)
        print("fpga:area={}, level={}, QoR={}".format(self.fX_area[index],self.fX_level[index],self.fX_obj[index]))
        self.engine.output(self.Sample_X[index])
        print("aig_area:",self.engine.area_0)
        
        # 输出结果质量，只用于测试
        with open(self.input_file + '.qor.txt', 'w') as file:
            file.write('level: ' + str(int(self.fX_level[index])) + '\n')
            file.write('area: ' + str(int(self.fX_area[index])) + '\n')
        self.plot()
        
    def time(self,t):
        with open(self.input_file + '.qor.txt', 'a') as file:
            file.write('time: ' + str(t) + '\n')
            
    def plot(self):
        plt.plot(list(range(len(self.fX_obj))),self.fX_obj)
        plt.title("QoR-iter plot")
        plt.xlabel("iter")# x轴名称 只能是英文
        plt.ylabel("QoR")# y轴名称 只能是英文
        # 显示图像
        plt.savefig(self.input_file + '.opt.png')
        

if __name__ == '__main__':     
    #input_file由用户输入   
    path = '../../benchmark_eda_elite/'
    DMA_comb = 'DMA_comb/DMA_comb.aig'
    b05 = 'b05_comb/b05_comb.aig'
    adder = 'adder/adder.aig'
    file_path = path+'b20_1_comb/b20_1_comb.aig'
    output_file = file_path + '.seq'    #或者None
    
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
        output_file = file_path + '.seq'    #或者None
        if len(sys.argv) > 2:
            output_file = sys.argv[2]
    
    print("--------------------------------------------------")        
    print(file_path)
            
    #time
    curr_time = time.time()

    #可设置参数
    seq_length = 7
    n_init = 5
    n_tol_eval = 10
    gp_train_steps = 300
    
    boils_exp = BOiLS(input_file=file_path, output_file=output_file)    
    area_size = boils_exp.engine.area_0
    
    if area_size < 1000:
        seq_length = 7
        n_init = 20
        n_tol_eval = 22
        gp_train_steps = 300
    elif area_size < 10000:
        seq_length = 8
        n_init = 20
        n_tol_eval = 30
        gp_train_steps = 500
    else:
        seq_length = 10
        n_init = 20
        n_tol_eval = 35
        gp_train_steps = 500
       
    boils_exp.run(n_tol_eval=n_tol_eval, n_init=n_init, seq_length=seq_length, gp_train_steps=gp_train_steps)

    end_time = time.time()
    time_cost = end_time - curr_time
    boils_exp.time(time_cost)
    print("time:",time_cost)